from collections import defaultdict, OrderedDict
import datetime
import logging
import functools

import pytz
from absl import (app, flags)

from coin.base.datetime_util import to_timestamp_int
from coin.support.accounting.logic.constants import COIN_LIST
from coin.support.proto_log.app.strat_log_replayer import _get_min_max_dt
from coin.support.pta.app.pta_dumper_runner import (calculate_pta_stats_from_strat_log,
                                                    gen_strategy_summary_message,
                                                    query_ticker)
from coin.support.pta.logic.pta_dumper import (dump_pta_stats,
                                               WeeklyStrategyPlotterConfig)
from coin.tool.strat_monitor.plot.strat_plotter import (
    get_week_range, StrategyPlotter)
from coin.tool.strat_monitor.util.calculator_util import adjust_pnl_from_config
from coin.tool.strat_monitor.util.monitor_util import (get_active_strategy,
                                                       get_strat_from_json_config,
                                                       get_pnl_adjust_info,
                                                       filter_empty_summary)
from coin.util.queue.config import KafkaConfig
from xunkemgmt_client.client.util.query_util import (
    query_accounts,
    query_exchange_apis,
    query_strategies)
from xunkemgmt_client.support.accounting.logic.constants import FIAT_LIST


FLAGS = flags.FLAGS


def _split_strategy_by_strategy_group(strats):
  strats_by_group = defaultdict(list)
  for strat in strats:
    strats_by_group[strat.strategy_group].append(strat)
  return strats_by_group


def main(_):
  year = FLAGS.year
  assert year >= 2019, year
  start_week = FLAGS.start_week
  assert start_week >= 0, start_week
  end_week = FLAGS.end_week
  assert end_week and end_week >= start_week, end_week
  strategy_config_filename = FLAGS.strategy_config_filename
  kafka_config_filename = FLAGS.kafka_config_filename
  strategy_group = FLAGS.strategy_group

  account_universe = query_accounts(as_proto=True)
  exchange_api_universe = query_exchange_apis(as_proto=True)
  strategy_universe = query_strategies(as_proto=True)

  strat_universe_by_group = _split_strategy_by_strategy_group(strategy_universe)
  strat_universe_by_group = OrderedDict(
      sorted(strat_universe_by_group.items(),
             key=lambda x: 1 if x[0].startswith('lm-') else 0,
             reverse=True))
  if strategy_config_filename is not None:
    strat_candidates = get_strat_from_json_config(strategy_config_filename)
  else:
    start_date, _ = get_week_range(year, start_week)
    threshold = to_timestamp_int(datetime.datetime.combine(
        start_date, datetime.time.min.replace(tzinfo=pytz.UTC)))
    strat_candidates = get_active_strategy(update_timestamp_threshold=threshold)

  for each_group, each_strats in strat_universe_by_group.items():
    logging.info('Running for strategy group: %s\n', each_group)
    all_strats = [elem.strategy_name for elem in each_strats]
    strat_list = set(strat_candidates).intersection(set(all_strats))
    if len(strat_list) == 0:
      continue
    if strategy_group == 'dmm':
      logging.info('Skip dmm strategies')
      continue
    if strategy_group is not None:
      if strategy_group != each_group:
        continue
    # TODO(chensili): Remove this.
    # strat_list = ['lm_agg_okex_btc_eos_eth_usdt']
    kafka_config = None
    if kafka_config_filename is not None:
      kafka_config = KafkaConfig.from_cmd_config(kafka_config_filename)

    for week in range(start_week, end_week):
      start_date, end_date = get_week_range(year, week)
      pnl_adj = get_pnl_adjust_info(
          FLAGS.pnl_adj_file, start_date, end_date + datetime.timedelta(days=1))
      logging.info('Running for week #%s: %s-%s ...\n' %
            (week, start_date.strftime('%Y%m%d'), end_date.strftime('%Y%m%d')))
      base_list = list(set(COIN_LIST + FIAT_LIST))
      yesterday = datetime.datetime.utcnow().date() - datetime.timedelta(days=1)
      query_date = min(yesterday, end_date)
      ticker = query_ticker(base_list, 'USD', trading_date=query_date)
      kline = None
      pnl_plot_func = None
      if FLAGS.plot_dir is not None:
        plotter_config = WeeklyStrategyPlotterConfig(year, week, FLAGS.plot_dir)
        plotter = StrategyPlotter(plotter_config, account_info=account_universe)
        pnl_plot_func = functools.partial(plotter.do_pnl_plot,
                                          axis_interval_hour=12)

      pta_summaries, details, num_fail = calculate_pta_stats_from_strat_log(
          start_dt=_get_min_max_dt(start_date)[0],
          end_dt=_get_min_max_dt(end_date)[1],
          strategy_list=strat_list,
          kafka_config=kafka_config,
          strat_log_root_dir=FLAGS.root_dir,
          max_workers=FLAGS.max_workers,
          ignore_error=FLAGS.ignore_error,
          pnl_adj=pnl_adj,
          realtime_pnl_adj=FLAGS.realtime_pnl_adj,
          account_universe=account_universe,
          exchange_api_universe=exchange_api_universe,
          strategy_universe=strategy_universe,
          ticker=ticker,
          kline=kline,
          pnl_plot_func=pnl_plot_func)
      pta_summaries = filter_empty_summary(pta_summaries)
      if not FLAGS.realtime_pnl_adj:
        pta_summaries, _ = adjust_pnl_from_config(
            pta_summaries, None, start_date, end_date, pnl_adj)
      summary_msg = gen_strategy_summary_message(end_date, pta_summaries)
      dump_pta_stats(summary_msg,
                     details,
                     ticker=ticker,
                     output_file=FLAGS.output_file)


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')

  flags.DEFINE_integer('year', None, 'Year')

  flags.DEFINE_integer('start_week', None, 'Start week.')

  flags.DEFINE_integer('end_week', None, 'End week.')

  flags.DEFINE_integer('max_workers', 24, 'Max number of workers.')

  flags.DEFINE_string('output_file', None, 'output_file')

  flags.DEFINE_string('plot_dir', None, 'plot dir')

  flags.DEFINE_boolean('ignore_error', False, 'Ignore error')

  flags.DEFINE_string('strategy_group', None, 'filter by strategy_group')

  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')

  flags.DEFINE_string('strategy_config_filename',
                      None,
                      'strategy_config_filename')

  flags.DEFINE_string('pnl_adj_file', None, 'pnl_adj_file')

  flags.DEFINE_boolean('realtime_pnl_adj', False, 'realtime pnl adj')

  flags.DEFINE_string('kline_db_config',
                      '../../coin_deploy/support_info/db_config_iosg/mysql_config_prod.json',
                      '')

  flags.DEFINE_string('kline_symbols_list',
                      '../../coin_deploy/support_info/kline_config/kline_symbols_list.csv',
                      'kline_symbols_list')

  flags.DEFINE_string('root_dir', None, '')

  app.run(main)
